from PIL import Image
import numpy as np  
import os
import pdb 

class ImgCroper:
    def __init__(self, img_folder_path, croped_save_path):
        try :
            os.path.exists(img_folder_path)
            os.path.exists(croped_save_path)
            self.img_folder_path = img_folder_path
            self.croped_save_path = croped_save_path
        except Exception as e:
            print("Path error",e)
    def has_white_pixels(self, img : "Image") -> bool:  
        # 将图像转换为NumPy数组  
        data = np.array(img)  
        
        # 检查是否存在白色像素（255, 255, 255）  
        if data.ndim == 3:  
            white_mask = (data[:, :, 0] == 255) & (data[:, :, 1] == 255) & (data[:, :, 2] == 255)    
        else:  
            # 如果图像不是RGB或RGBA，则无法检查白色像素  
            raise TypeError("文件类型不支持的操作")
        return np.any(white_mask)
        
    def make_white_transparent(self, img : 'Image'):  
        # 检查图像模式，并根据需要添加alpha通道  
        if img.mode == 'RGB':  
            # 创建一个新的图像，添加一个全为255的alpha通道  
            img = img.convert('RGBA')  
        elif img.mode != 'RGBA':  
            # 如果图像不是RGBA模式，则抛出错误  
            raise ValueError("Image is not an RGBA image.")  
        # 将白色转换为透明  
        data = np.array(img)  
        # 假设白色是(255, 255, 255)，我们将其alpha值设置为0  
        # 检查数组是否是3维的，并且最后一个维度是4（对于RGBA图像）  
        if data.ndim == 3 and data.shape[2] == 4:  
            # 将白色转换为透明  
            white_mask = (data[:, :, 0] == 255) & (data[:, :, 1] == 255) & (data[:, :, 2] == 255)  
            data[white_mask, 3] = 0  # 将白色部分的alpha值设置为0  
        else:  
            print("Error: Image is not an RGBA image.")  
            return None   
        # 将修改后的数组转换回PIL.Image对象  
        img = Image.fromarray(data)  
        return img
    def find_bounding_box(self,img : 'Image'): 
        try:  
            width, height = img.size  
        except Exception as e: 
            print("img:", img)   
            raise               # 重新抛出异常，以便上层代码可以处理它。  
        
        # 初始化边界变量  
        left, upper = width, height  
        right, lower = 0, 0  
        
        # 辅助函数，用于检查像素是否不透明  
        def is_opaque(pixel):  
            return pixel[3] == 255  
        
        # 遍历图像像素  
        pixels = img.load()  
        for x in range(width):  
            for y in range(height):  
                pixel = pixels[x, y]  
                if is_opaque(pixel):  
                    if x < left:  
                        left = x  
                    if x > right:  
                        right = x  
                    if y < upper:  
                        upper = y  
                    if y > lower:  
                        lower = y  
        
        # 确定边界框  
        box = (left, upper, right + 1, lower + 1)  
        return box  

    def crop_save_images(self) -> None:
        for filename in os.listdir(self.img_folder_path):
            if filename.endswith(".png"):
                img_path = os.path.join(self.img_folder_path, filename)
                img_output_path = os.path.join(self.croped_save_path, filename)
                try:
                    img = Image.open(img_path)
                    if self.has_white_pixels(img):
                        print("白色图片：",img_path)
                        img = self.make_white_transparent(img)
                    box = self.find_bounding_box(img)
                    cropped_image = img.crop(box)
                    cropped_image.save(img_output_path)
                    # assert img is None,"图片读取错误"
                except Exception as e:
                    print("图片读取错误",e)
                    print("图片路径：",img_path)

def test_make_white_transparent():
    margin_white_img_path = '/root/MyCode/PNG_image_crop/data/5690.png'
    output_folder = r'/root/MyCode/PNG_image_crop/output'
    output_path = os.path.join(output_folder, 'cropped_image.png')  
    make_white_transparent(margin_white_img_path,output_path)

def test_crop():           
    input_folder = r'/root/MyCode/PNG_image_crop/data'
    output_folder = r'/root/MyCode/PNG_image_crop/output'
    crop_text_images(input_folder, output_folder)

def test_confirm_box():
    # 使用函数  
    image_path = '/root/MyCode/PNG_image_crop/data/cropped_image.png' 
    # image_path_white = '/root/MyCode/PNG_image_crop/data/5690.png'
    output_folder = r'/root/MyCode/PNG_image_crop/output'
    output_path = os.path.join(output_folder, 'cropped_image.png')  
    box = find_bounding_box(image_path)  
    print("Bounding box:", box)  
    
def test_whiteimg_is_white():
    margin_white_img_path = '/root/MyCode/PNG_image_crop/data/5690.png'
    margin_white_img = Image.open(margin_white_img_path)
    instance = ImgCroper('/root/MyCode/PNG_image_crop/data','/root/MyCode/PNG_image_crop/output')
    if instance.has_white_pixels(margin_white_img):
        return True

def test_transparentimg_is_not_white():
    without_white_img_path = '/root/MyCode/PNG_image_crop/data/15163.png'
    without_white_img = Image.open(without_white_img_path)
    instance = ImgCroper('/root/MyCode/PNG_image_crop/data','/root/MyCode/PNG_image_crop/output')
    if instance.has_white_pixels(without_white_img) == False:
        return True


def test():
    assert test_whiteimg_is_white() == True
    assert test_transparentimg_is_not_white() == True
    #test_make_white_transparent()
# test()


def test_crop_save_imgs():
    images_path = '/root/MyCode/PNG_image_crop/data'
    save_path = '/root/MyCode/PNG_image_crop/output'
    instance = ImgCroper(images_path,save_path)
    instance.crop_save_images()
# test_crop_save_imgs()

def test_crop_save_imgs():  
    images_path = input("请输入图片所在目录的路径: ")  
    save_path = input("请输入裁剪后图片保存的目录路径: ")  
    instance = ImgCroper(images_path, save_path)  
    instance.crop_save_images()  
test_crop_save_imgs()
